package com.mdd.front.factory.chatAi;


import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Strings;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.theokanning.openai.OpenAiApi;
import com.theokanning.openai.completion.chat.ChatCompletionChoice;
import com.theokanning.openai.completion.chat.ChatCompletionChunk;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatMessage;
import com.theokanning.openai.service.OpenAiService;
import io.github.asleepyfish.config.ChatGPTProperties;
import io.github.asleepyfish.enums.ChatGPTErrorEnum;
import io.github.asleepyfish.enums.RoleEnum;
import io.github.asleepyfish.exception.ChatGPTException;
import io.github.asleepyfish.service.OpenAiProxyService;
import okhttp3.ConnectionPool;
import okhttp3.OkHttpClient;
import org.springframework.transaction.annotation.Transactional;
import retrofit2.Retrofit;

import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.nio.charset.Charset;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

public class ChatBaseOpenAiProxyService extends OpenAiService {


    private static final String BASE_URL = "https://api.openai.com/";

    private static final Random RANDOM = new Random();
    private final Cache<String, LinkedList<ChatMessage>> cache;

    private final ChatGPTProperties chatGPTProperties;
    private final OkHttpClient client;

    public void createStreamChatCompletion(String content, String userSessionId, OutputStream os,StringBuilder answer) {
        createStreamChatCompletion(content, userSessionId, chatGPTProperties.getChatModel(), os,answer);
    }

    public void createStreamChatCompletion(String content, String userSessionId, String model, OutputStream os,StringBuilder answer) {
        createStreamChatCompletion(RoleEnum.USER.getRoleName(), content, userSessionId, model, 1.0D, 1.0D, os,answer);
    }


    public void createStreamChatCompletion(String role, String content, String userSessionId, String model, Double temperature, Double topP, OutputStream os,StringBuilder answer) {
        createStreamChatCompletion(ChatCompletionRequest.builder()
                .model(model)
                .messages(Collections.singletonList(new ChatMessage(role, content)))
                .user(userSessionId)
                .temperature(temperature)
                .topP(topP)
                .stream(true)
                .build(),
                os,
                answer);
    }

    public void createStreamChatCompletion(ChatCompletionRequest chatCompletionRequest, OutputStream os,StringBuilder answer) {
        chatCompletionRequest.setStream(true);
        chatCompletionRequest.setN(1);
        String user = chatCompletionRequest.getUser();
        LinkedList<ChatMessage> contextInfo = new LinkedList<>();
        try {
            contextInfo = cache.get(user, LinkedList::new);
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        contextInfo.addAll(chatCompletionRequest.getMessages());
        chatCompletionRequest.setMessages(contextInfo);
        List<ChatCompletionChunk> chunks = new ArrayList<>();
        for (int i = 0; i < chatGPTProperties.getRetries(); i++) {
            try {
                // 避免频繁要求，随机睡眠0.5s ~ 0.7s
                if (i > 0) {
                    randomSleep();
                }
                super.streamChatCompletion(chatCompletionRequest).doOnError(Throwable::printStackTrace).blockingForEach(chunk -> {
                    chunk.getChoices().stream().map(choice -> choice.getMessage().getContent())
                            .filter(Objects::nonNull).findFirst().ifPresent(o -> {
                                try {
                                    os.write(o.getBytes(Charset.defaultCharset()));
                                    os.flush();
                                    answer.append(o);
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                            });
                    chunks.add(chunk);
                });
                // 如果最后一行代码是正确的，我们可以简单地打破循环
                break;
            } catch (Exception e) {
                String message = e.getMessage();
                int size = Objects.requireNonNull(cache.getIfPresent(user)).size();
                for (int j = 0; j < size / 2; j++) {
                    Objects.requireNonNull(cache.getIfPresent(user)).removeFirst();
                }
                chatCompletionRequest.setMessages(cache.getIfPresent(user));
                if (i == chatGPTProperties.getRetries() - 1) {
                    e.printStackTrace();
                    throw new ChatGPTException(ChatGPTErrorEnum.FAILED_TO_GENERATE_ANSWER, message);
                }
            }
        }
        LinkedList<ChatMessage> chatMessages = new LinkedList<>();
        try {
            chatMessages = cache.get(user, LinkedList::new);
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        chatMessages.add(new ChatMessage(RoleEnum.ASSISTANT.getRoleName(), chunks.stream()
                .flatMap(chunk -> chunk.getChoices().stream())
                .map(ChatCompletionChoice::getMessage)
                .map(ChatMessage::getContent)
                .filter(Objects::nonNull)
                .collect(Collectors.joining())));
    }

    public ChatBaseOpenAiProxyService(ChatGPTProperties chatGPTProperties, Duration timeout) {
        super(buildApi(chatGPTProperties.getToken(), timeout, chatGPTProperties.getProxyHost(), chatGPTProperties.getProxyPort()),
                defaultClient(chatGPTProperties.getToken(), timeout, chatGPTProperties.getProxyHost(), chatGPTProperties.getProxyPort()).dispatcher().executorService());
        this.chatGPTProperties = chatGPTProperties;
        this.cache = chatGPTProperties.getSessionExpirationTime() == null ? CacheBuilder.newBuilder().build() :
                CacheBuilder.newBuilder().expireAfterAccess(chatGPTProperties.getSessionExpirationTime(), TimeUnit.MINUTES).build();
        this.client = OpenAiProxyService.defaultClient(chatGPTProperties.getToken(), Duration.ZERO, chatGPTProperties.getProxyHost(), chatGPTProperties.getProxyPort());
    }

    public static OpenAiApi buildApi(String token, Duration timeout, String proxyHost, int proxyPort) {
        ObjectMapper mapper = defaultObjectMapper();
        OkHttpClient client = defaultClient(token, timeout, proxyHost, proxyPort);
        Retrofit retrofit = defaultRetrofit(client, mapper);
        return retrofit.create(OpenAiApi.class);
    }

    private void randomSleep() throws InterruptedException {
        Thread.sleep(500 + RANDOM.nextInt(200));
    }

    public static OkHttpClient defaultClient(String token, Duration timeout, String proxyHost, int proxyPort) {
        if (Strings.isNullOrEmpty(proxyHost)) {
            return new OkHttpClient.Builder()
                    .addInterceptor(new ChatBaseAuthenticationInterceptor(token))
                    .connectionPool(new ConnectionPool(20, 5, TimeUnit.SECONDS))
                    .readTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS)
                    .build();
        }
        // 创建代理对象
        Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxyHost, proxyPort));
        return OpenAiService.defaultClient(token, timeout).newBuilder().proxy(proxy).build();
    }


}
